package com.nervenets.general.utils.builder;

import org.springframework.amqp.utils.MapBuilder;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;

public class ThreadArgumentsBuilder<K, V> extends MapBuilder<ThreadArgumentsBuilder<K, V>, K, V> {
    private final ConcurrentHashMap<K, Boolean> threadMap = new ConcurrentHashMap<>();
    private final List<Thread> threads = new ArrayList<>();
    private final Object queueLock = new Object();

    @Override
    public ThreadArgumentsBuilder<K, V> put(K key, V value) {
        final ThreadArgumentsBuilder<K, V> _this = super.put(key, value);
        threadMap.put(key, true);
        return _this;
    }

    public ThreadArgumentsBuilder<K, V> putIf(boolean condition, K key, Supplier<V> supplier) {
        if (condition) {
            final RequestAttributes attributes = RequestContextHolder.getRequestAttributes();
            threadMap.put(key, false);
            threads.add(new Thread(() -> {
                try {
                    RequestContextHolder.setRequestAttributes(attributes);
                    put(key, supplier.get());
                } finally {
                    RequestContextHolder.resetRequestAttributes();
                    threadMap.put(key, true);
                    synchronized (queueLock) {
                        queueLock.notifyAll();
                    }
                }
            }));
        }
        return this._this();
    }

    public ThreadArgumentsBuilder<K, V> put(K key, Supplier<V> supplier) {
        return this.putIf(true, key, supplier);
    }

    @Override
    public Map<K, V> get() {
        threads.forEach(Thread::start);
        synchronized (queueLock) {
            while (threadMap.containsValue(false)) {
                try {
                    queueLock.wait();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
            threads.clear();
            threadMap.clear();
            return super.get();
        }
    }
}
